﻿using System;
using System.Collections.Generic;
using System.Text;
using System.Windows.Forms;
using System.ComponentModel;

namespace Axe.Windows.Forms
{
    public class AutoExpandNode : TreeNode
    {
        TreeNode stub = null;

        public AutoExpandNode(bool hasChilds)
        {
            if (hasChilds)
            {
                stub = new TreeNode();
                Nodes.Add(stub);
                Collapse();

                CanContainChilds = true;
                NeedPopulate = true;
            }
        }

        internal void RemoveStubNode()
        {
            if(stub!=null)
                Nodes.Remove(stub);
        }

        bool m_CanContainChilds = false;
        bool m_IsPopulating = false;
        bool m_NeedPopulate = false;

        public bool NeedPopulate
        {
            get { return m_NeedPopulate; }
            set { m_NeedPopulate = value; }
        }

        string m_PopulatingTextFormatString = "{0} (loading...)";
        string savedNodeText = null;
        object m_PopulateNodeKey = null;

        public virtual object PopulateNodeKey
        {
            get { return m_PopulateNodeKey; }
            set { m_PopulateNodeKey = value; }
        }

        public string PopulatingTextFormatString
        {
            get { return m_PopulatingTextFormatString; }
            set { m_PopulatingTextFormatString = value; }
        }

        public bool IsPopulating
        {
            get { return m_IsPopulating; }
            set { m_IsPopulating = value; }
        }

        public bool CanContainChilds
        {
            get { return m_CanContainChilds; }
            set { m_CanContainChilds = value; }
        }

        public virtual void UpdateNodeText()
        {
            if (IsPopulating)
            {
                savedNodeText = this.Text;
                this.Text = String.Format(PopulatingTextFormatString, Text);
            }
            else
            {
                this.Text = savedNodeText;
            }
        }

        public void ResetExpand()
        {
            NeedPopulate = true;
            Nodes.Clear();
            Nodes.Add(stub);
            Collapse();
        }
    }

    public class PopulatingEventArgs : TreeViewCancelEventArgs
    {
        bool m_ClearExisted = false;
        IEnumerable<TreeNode> m_Results;
        object m_PopulateKey;

        public object PopulateKey
        {
            get { return m_PopulateKey; }
            set { m_PopulateKey = value; }
        }

        public bool ClearExisted
        {
            get { return m_ClearExisted; }
            set { m_ClearExisted = value; }
        }

        public IEnumerable<TreeNode> Results
        {
          get { return m_Results; }
          set { m_Results = value; }
        }

        public PopulatingEventArgs(AutoExpandNode node)
            :base(node, false, TreeViewAction.Expand)
        {
            PopulateKey = node.PopulateNodeKey;
        }
    }

    public class AutoPopulateTreeView : TreeView
    {
        readonly Dictionary<object, AutoExpandNode> ExpandNodes = new Dictionary<object, AutoExpandNode>();

        Dictionary<AutoExpandNode, PopulatingEventArgs> asyncPopulate = new Dictionary<AutoExpandNode, PopulatingEventArgs>();
        bool useAsyncPopulate = false;

        [DefaultValue(false)]
        public bool UseAsyncPopulate
        {
            get { return useAsyncPopulate; }
            set { useAsyncPopulate = value; }
        }

        protected override void OnBeforeExpand(TreeViewCancelEventArgs e)
        {
            if (e.Node is AutoExpandNode)
            {
                OnNodePopulate(e.Node as AutoExpandNode);
            }

            base.OnBeforeExpand(e);
        }


        protected virtual void OnNodePopulate(AutoExpandNode node)
        {
            node.RemoveStubNode();

            if (node.NeedPopulate == false)
                return;

            if (NodePopulate != null)
            {
                node.IsPopulating = true;
                node.UpdateNodeText();

                PopulatingEventArgs args = new PopulatingEventArgs(node);

                if (useAsyncPopulate)
                {
                    lock (asyncPopulate)
                    {
                        if (!asyncPopulate.ContainsKey(node))
                        {
                            asyncPopulate[node] = args;
                        }
                        else
                            return;
                    }

                    NodePopulate.BeginInvoke(this, args, PopulateAsyncCallback, args);
                }
                else
                {
                    NodePopulate(this, args);
                    EndPopulate(args);
                }
            }
        }

        private void PopulateAsyncCallback(IAsyncResult res)
        {
            PopulatingEventArgs args = (PopulatingEventArgs)res.AsyncState;

            lock (asyncPopulate)
            {
                asyncPopulate.Remove(args.Node as AutoExpandNode);
            }

            MethodInvoker mi = new MethodInvoker(delegate() { EndPopulate(args); });
            this.Invoke(mi);
        }

        private void EndPopulate(PopulatingEventArgs args)
        {
            AutoExpandNode node = args.Node as AutoExpandNode;

            if (!args.Cancel)
            {
                if (args.ClearExisted)
                {
                    node.Nodes.Clear();
                }

                if (args.Results != null)
                {
                    foreach (TreeNode child in args.Results)
                    {
                        node.Nodes.Add(child);
                    }
                }

                node.NeedPopulate = false;
            }

            node.IsPopulating = false;
            node.UpdateNodeText();
            node.Expand();
            UpdateNodesHashtable();
        }

        public event EventHandler<PopulatingEventArgs> NodePopulate;

        private void UpdateNodesHashtable(TreeNodeCollection collection)
        {
            if (collection == null)
                return;

            foreach (TreeNode node in collection)
            {
                if (node is AutoExpandNode)
                {
                    AutoExpandNode expNode = (node as AutoExpandNode);
                    ExpandNodes[expNode.PopulateNodeKey] = expNode;
                }

                UpdateNodesHashtable(node.Nodes);
            }
        }

        public void UpdateNodesHashtable()
        {
            ExpandNodes.Clear();
            UpdateNodesHashtable(this.Nodes);
        }

        public AutoExpandNode GetPopulateNode(object key)
        {
            if (ExpandNodes.ContainsKey(key))
            {
                return ExpandNodes[key];
            }

            return null;
        }
    }
}
